import torch
from torchvision.models import resnet50, resnet152

if __name__ == '__main__':
    # 虽然这里设置cuda:0，但实际使用的是1号gpu
    device = torch.device('cuda:0' if torch.cuda.is_available else 'cpu')
    print(f'当前设备为：{torch.cuda.current_device()}')
    model = resnet152(num_classes=10)
    model.to(device)
    # 使用res152做1000次前向推断，batch-size设置为16
    for i in range(1000):
        X = torch.randn(16, 3, 224, 224).to(device)
        y = model(X)
        print(f'id:{i + 1:3d}:{y}')